﻿/****************************************************************************
** file: RotateMatchTemplate.h
** brief: 基于 OpenCV matchTemplete 实现的模板匹配算法，支持图像旋转匹配。
** Copyright (C) LiYuan
** Author: LiYuan
** E-Mail: 18069211#qq(.)com
** Version 0.1.1
** Last modified: 2024.03.27
** Modified By: LiYuan
****************************************************************************/


#ifndef ROTATEMATCHTEMPLATE_H
#define ROTATEMATCHTEMPLATE_H

#include "ImageRotate.h"
#include <opencv2/core.hpp>
#include <vector>

namespace Qly {

/**
 * @brief drawRotatedRect 在图像上画出旋转的矩形
 * @param [inout] image
 * @param [in] rrt 旋转的举行
 * @param color 颜色
 * @param thickness 线宽
 */
bool drawRotatedRect(cv::Mat &image, cv::RotatedRect &rrt,
                     cv::Scalar color = cv::Scalar(255, 255, 255), int thickness = 1);


/**
 * @brief The TemplateMatch class 这个类主要解决图像旋转匹配问题。 如果不考虑旋转那么 match 函数就够用了。
 * TODO: 还没有完整实现亚像素精度。里面只有角度插值算法。还没写位置插值算法。
 */
class RotateTemplateMatch
{
public:
    enum class MODEL_SIZE_MODE {FIX_MODEL_SIZE, AUTO_MODEL_SIZE, FULL_MODEL_SIZE};

    /**
     * @brief setMethod 设置寻找最佳角度的方法
     * @param method  支持 TM_SQDIFF， TM_SQDIFF_NORMED， TM_CCORR， TM_CCOEFF， TM_CCORR_NORMED， TM_CCOEFF_NORMED
     */
    void setMethod( int method = cv::TM_CCORR ) {m_matchMethod = method;}
    void setSizeMode(MODEL_SIZE_MODE mode) {m_sizeMode = mode;}
    void setSubPixMode(bool enable) {m_enableSubPix = enable;}

    bool setModelByAngleCount(const cv::Mat &model, double angleMin, double angleMax, size_t angleCount);
    bool setModelByAngleCount(const cv::Mat &model, const cv::Mat &mask, double angleMin, double angleMax, size_t angleCount);

    /**
     * @brief setModelByDetAngle 生成各个角度的 model
     * @param model 标准 model
     * @param angleMin 起始角度，以度为单位，要求大于 -360 度
     * @param angleMax 终止角度，以度为单位，要求小于 360 度
     * @param detAngle 角度间隔
     * @return
     */
    bool setModelByDetAngle(const cv::Mat &model, double angleMin, double angleMax, double detAngle);
    bool setModelByDetAngle(const cv::Mat &model, const cv::Mat &mask, double angleMin, double angleMax, double detAngle);

    /**
     * @brief singleMatch 针对特定角度做一次匹配。 这个函数可以开启亚像素精度的位置计算。
     * @param image 输入的图像
     * @param angleIndex 角度 index
     * @param rrt 返回匹配到的位置
     * @param score 返回匹配的分数
     * @return
     */
    bool singleMatch(const cv::Mat &image, int angleIndex, cv::RotatedRect &rrt, double &score);
    bool singleMatch(const cv::Mat &image, int angleIndex, cv::Rect &rect, double &score);

    /**
     * @brief rotatedMatch 旋转匹配，对各个角度都进行一次匹配，找出最佳角度。这个函数可以开启亚像素精度的位置计算。
     * @param image 输入的图像
     * @param rrt 返回匹配到的位置
     * @param score 返回匹配的分数
     * @return
     */
    bool rotatedMatch(const cv::Mat &image, cv::RotatedRect &rrt, double &score);
    /**
     * @brief rotatedMatch 旋转匹配，对各个角度都进行一次匹配，找出最佳角度
     * @param image 输入的图像
     * @param rect 返回匹配到的位置
     * @param index 返回匹配到的 angle index.
     * @param score 返回匹配的分数
     * @return
     */
    bool rotatedMatch(const cv::Mat &image, cv::Rect &rect, size_t &index, double &score);

    size_t angleCount() {return m_angleCount;}

    /**
     * @brief rectSize 返回多个 model 的最大长、宽尺寸。
     * @param indexMin 角度范围
     * @param indexMax 角度范围
     * @param width 返回宽度值
     * @param height 返回高度值
     * @return
     */
    bool rectSize(size_t indexMin, size_t indexMax, int &width, int &height); //返回这个范围内的 model 的尺寸

    /**
     * @brief rotatedModel 获取特定角度的 model
     * @param angleIndex 角度 index
     * @param model 返回获取的 model
     * @return false 表示失败，true 表示成功
     */
    bool rotatedModel(size_t angleIndex, cv::Mat &model);
    bool rotatedMask(size_t angleIndex, cv::Mat &mask);
protected:

    // 内部函数，用来判断当前 score 是否比之前的更优
    bool better(double oldGrade, double newGrade, int method);

    // 内部函数，角度插值，计算最佳角度
    bool angleInterpolate(size_t index, double &angle, double& score);

    // 内部函数，被 setModelByxxx 函数调用。用来生成真正的旋转的 models
    //bool buildRotatedModels(const cv::Mat &model, const cv::Mat &mask);
    bool buildRotatedModels(const cv::Mat &model, const cv::Mat &mask);

    cv::Mat m_model;
    cv::Mat m_mask;
    cv::Size m_modelSize;
    std::vector<cv::Mat> m_rotatedModels;
    std::vector<cv::Mat> m_rotatedMasks;
    std::vector<double> m_angles;

    std::vector<double> m_scores; // 用来做角度插值
    cv::Mat m_result; // 用来做位置插值


    double m_angleMin;
    double m_angleMax;
    double m_angleDet;
    size_t m_angleCount;

    MODEL_SIZE_MODE m_sizeMode = MODEL_SIZE_MODE::FIX_MODEL_SIZE;
    int m_matchMethod = cv::TM_CCORR_NORMED;
    bool m_enableSubPix = false;
};
}


#endif // ROTATEMATCHTEMPLATE_H
